from transformers import AutoTokenizer,AutoModelForSequenceClassification
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from torch.utils.data import random_split
import torch
from torch.optim import Adam
import evaluate
# data = pd.read_csv("D:\datasets\ChnSentiCorp\ChnSentiCorp_htl_all.csv")
# data = data.dropna() #删除空行

class MyDataset(Dataset):
    def __init__(self):
        super().__init__()
        self.data = pd.read_csv("D:\datasets\ChnSentiCorp\ChnSentiCorp_htl_all.csv")
        self.data = self.data.dropna()
    def __getitem__(self, index):
        return self.data.iloc[index]["review"],self.data.iloc[index]["label"]
    def __len__(self):
        return len(self.data)

dataset = MyDataset()
# for i in range(5):
#     print(dataset[i])

#划分数据集
trainset,validset = random_split(dataset,[0.8,0.2])
print(len(trainset),len(validset))

#创建数据加载器
tokenizer = AutoTokenizer.from_pretrained(r"D:\models\rbt3")

def collate_fn(batch):
    texts,labels = [],[]
    for item in batch:
        texts.append(item[0])
        labels.append(item[1])
    inputs = tokenizer(texts,max_length=128,truncation=True,padding="max_length",return_tensors="pt")
    inputs['labels'] = torch.tensor(labels)
    return inputs

trainloader = DataLoader(trainset,batch_size=100,shuffle=True,collate_fn=collate_fn)
validloader = DataLoader(validset,batch_size=50,shuffle=True,collate_fn=collate_fn)

# print(next(enumerate(trainloader))[1]) #打印训练数据加载器中第一个批次的数据

#创建模型
model = AutoModelForSequenceClassification.from_pretrained(r"D:\models\rbt3")
if torch.cuda.is_available():
    model = model.cuda()
optimizer = Adam(model.parameters(),lr=1e-5)

#训练+验证
cls_metric = evaluate.combine([r"D:\Code\sshcode\HuggingFace\evaluate-main\metrics\accuracy",
                               r"D:\Code\sshcode\HuggingFace\evaluate-main\metrics\f1",
                               r"D:\Code\sshcode\HuggingFace\evaluate-main\metrics\recall",
                               r"D:\Code\sshcode\HuggingFace\evaluate-main\metrics\precision"])
def evaluate():
    model.eval()
    # acc_num = 0
    with torch.inference_mode():
        for batch in validloader:
            if torch.cuda.is_available():
                batch = {k:v.cuda() for k,v in batch.items()}
                output = model(**batch)
                pred = torch.argmax(output.logits,dim=-1)
                # acc_num += (batch['labels'].long() == pred.long()).float().sum()
                cls_metric.add_batch(references=batch['labels'].long(),predictions=pred.long())
    return cls_metric.compute()

def train(epoch=3):
    for en in range(epoch):
        train_sum_loss = 0
        model.train()
        for batch in trainloader:
            if torch.cuda.is_available():
                batch = {k:v.cuda() for k,v in batch.items()}
            output = model(**batch)
            optimizer.zero_grad()
            output.loss.backward()
            optimizer.step()

            train_sum_loss += output.loss.item()
        acc = evaluate()
        train_avg_loss = train_sum_loss/ len(trainloader)
        print("epoch:",en,"train_avg_loss:",train_avg_loss,"acc:",acc)

#模型预测
# sen  = "这家酒店不错！饭很好吃"
# id2_label = {0:"差评！",1:"好评"}
# model.eval()
# with torch.inference_mode():
#     inputs = tokenizer(sen,max_length=128,truncation=True,padding="max_length",return_tensors="pt")
#     inputs = {k:v.cuda() for k,v in inputs.items()}
#     output = model(**inputs)
#     logits = torch.softmax(output.logits,dim=-1)
#     pred = torch.argmax(logits,dim=-1)
#     print(f"输出结果:{id2_label.get(pred.item())}")

# from transformers import pipeline
# pipe = pipeline("text-classification",model=model,tokenizer=tokenizer,device=0)
# res = pipe(sen)
# print(res)

if __name__ == '__main__':
    train()